from wrf import (getvar, interplevel)
from Func_List_Files import list_ncfiles
from Func_Create_File import create_file
from math import (atan, cos, sin, sqrt)
from netCDF4 import Dataset
import numpy as np
import math
import csv
import os



# Define all hurricane's settings 
#---------------------------------------------------------------------------------------

Input_Dir  = '.../Non_Hurricane_Case/'           # Provide input directory
Output_Dir = '.../Post_Processing/Snapshot/'     # Provide output directory
HNS        = ['August_19']
HR         = ['168']
GSS        = ['32km']
TMS        = ['NoTurb']
CLZS       = ['CLZ0p0001', 'CLZ100p0', 'Default']
Time_Idx   = 5

# Extract Azimuthal wind velocity snapshot
#---------------------------------------------------------------------------------------
idx = 0
for HN in HNS:
    for GS in GSS:
        for TM in TMS:
            clz_count = 0
            for CLZ in CLZS:
                ncfiles           = []
                Hurricane_Setting = HN + '_' + HR[idx] + 'Hours_' + GS + '_' + TM + '_' + CLZ
                ncfiles           = list_ncfiles (Input_Dir + Hurricane_Setting, ncfiles)
                
                for ncfile in ncfiles[:-1]:
                    os.chdir(Input_Dir + Hurricane_Setting)
                    ncfile_name = ncfile
                    ncfile = Dataset(ncfile)
                    lat     = np.array(getvar(ncfile, "XLAT", Time_Idx))[:,0]
                    lon     = np.array(getvar(ncfile, "XLONG", Time_Idx))[0,:]
                    WSPD    = np.zeros((len(lat), len(lon)))
                    U10     = np.array(getvar(ncfile, "U10", Time_Idx))

# Export Azimuthal wind velocity snapshot
#---------------------------------------------------------------------------------------
                    create_file (Output_Dir, Hurricane_Setting)
                    file_name = ncfile_name + '_' + str(Time_Idx)
                    MyFile=open("%s.csv" %file_name,'w')
                    MyFile.write(',')
                    for i in range (len(lon)):
                        MyFile.write(str(lon[i]) + ',')
                    MyFile.write('\n')

                    for n in range(len(lat)):
                        MyFile.write(str(lat[n]) + ',')
                        for m in range (len(lon)):
                            MyFile.write(str(U10[n,m]) + ',')
                        MyFile.write('\n')

                clz_count += 1
    idx+=1        
